# -*- coding: utf-8 -*-
"""
直接导入数据（绕过配置文件问题）
通过临时重命名配置文件来避免编码错误
"""

import os
import sys
from pathlib import Path

# 设置环境变量
os.environ['PGCLIENTENCODING'] = 'UTF8'
os.environ['PYTHONIOENCODING'] = 'utf-8'

# 禁用所有可能的配置文件读取
os.environ['PGPASSFILE'] = 'nul'  # Windows 上的空设备
os.environ['PGSYSCONFDIR'] = ''
os.environ['PGSERVICEFILE'] = ''

project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

def main():
    print("=" * 60)
    print("直接导入数据（绕过方案）")
    print("=" * 60)
    
    import psycopg2
    from psycopg2.extras import Json
    from D_db.core.tokenizer import Tokenizer
    from D_db.io.excel_reader import ExcelReader
    
    # 直接在代码中指定所有连接参数
    conn_params = {
        'host': 'localhost',
        'port': 5432,
        'dbname': 'QA_knowledge',
        'user': 'postgres',
        'password': '123456',
        'options': '-c client_encoding=UTF8',
        'sslmode': 'disable',  # 禁用 SSL
        'connect_timeout': 10
    }
    
    print("\n步骤1: 连接数据库...")
    try:
        # 使用最简单的连接方式
        conn = psycopg2.connect(**conn_params)
        print("✓ 数据库连接成功")
    except Exception as e:
        print(f"✗ 连接失败: {e}")
        print("\n请手动在 psql 中创建数据库：")
        print("  psql -U postgres -c \"DROP DATABASE IF EXISTS QA_knowledge;\"")
        print("  psql -U postgres -c \"CREATE DATABASE QA_knowledge ENCODING 'UTF8';\"")
        return
    
    try:
        # 步骤2: 创建表
        print("\n步骤2: 创建数据库表...")
        cursor = conn.cursor()
        
        # 删除旧表
        cursor.execute("DROP TABLE IF EXISTS qa_knowledge;")
        
        # 创建新表
        cursor.execute("""
            CREATE TABLE qa_knowledge (
                question TEXT PRIMARY KEY,
                question_tokens JSONB,
                answer JSONB,
                context TEXT,
                category_id VARCHAR(50),
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
            );
        """)
        
        # 创建索引
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_qa_question_tokens_gin 
            ON qa_knowledge USING GIN (question_tokens);
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_qa_tokens_array 
            ON qa_knowledge USING GIN ((question_tokens->'tokens'));
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_qa_category_id 
            ON qa_knowledge (category_id);
        """)
        
        conn.commit()
        print("✓ 数据库表创建成功")
        
        # 步骤3: 读取Excel
        print("\n步骤3: 读取Excel文件...")
        excel_file = project_root / "C_filter" / "QA_final" / "大伙房水库调度规程.xlsx"
        
        if not excel_file.exists():
            print(f"✗ 文件不存在: {excel_file}")
            return
        
        print(f"  文件: {excel_file.name}")
        
        excel_reader = ExcelReader()
        df = excel_reader.read_file(str(excel_file))
        qa_pairs = excel_reader.parse_qa_data(df)
        
        if not qa_pairs:
            print("✗ 没有找到有效的QA数据")
            return
        
        print(f"✓ 找到 {len(qa_pairs)} 个问答对")
        
        # 步骤4: 分词
        print("\n步骤4: 对问题进行分词...")
        tokenizer = Tokenizer()
        processed_qa_pairs = []
        
        for i, (question, answer, context) in enumerate(qa_pairs, 1):
            if i % 500 == 0:
                print(f"  已处理 {i}/{len(qa_pairs)} 个问答对...")
            
            token_data = tokenizer.tokenize_question(question)
            processed_qa_pairs.append((question, answer, token_data, context))
        
        print(f"✓ 分词完成")
        
        # 步骤5: 插入数据库
        print("\n步骤5: 插入数据库...")
        
        insert_sql = """
        INSERT INTO qa_knowledge (question, question_tokens, answer, context, category_id)
        VALUES (%s, %s, %s, %s, %s)
        ON CONFLICT (question) DO UPDATE SET
            question_tokens = EXCLUDED.question_tokens,
            answer = EXCLUDED.answer,
            context = EXCLUDED.context,
            category_id = EXCLUDED.category_id,
            updated_at = CURRENT_TIMESTAMP;
        """
        
        batch_size = 100
        total = len(processed_qa_pairs)
        
        for i in range(0, total, batch_size):
            batch = processed_qa_pairs[i:i+batch_size]
            
            for question, answer, token_data, context in batch:
                answer_json = {"answer": answer}
                cursor.execute(insert_sql, (
                    question,
                    Json(token_data),
                    Json(answer_json),
                    context,
                    "21100150"  # 大伙房水库
                ))
            
            conn.commit()
            print(f"  已插入 {min(i+batch_size, total)}/{total} 条记录...")
        
        print(f"✓ 所有数据插入完成")
        
        # 验证
        cursor.execute("SELECT COUNT(*) FROM qa_knowledge;")
        count = cursor.fetchone()[0]
        
        print("\n" + "=" * 60)
        print("✓ 导入完成！")
        print("=" * 60)
        print(f"数据库: QA_knowledge")
        print(f"表名: qa_knowledge")
        print(f"总记录数: {count}")
        
    except Exception as e:
        print(f"\n✗ 执行失败: {e}")
        import traceback
        traceback.print_exc()
        conn.rollback()
    finally:
        cursor.close()
        conn.close()
        print("\n数据库连接已关闭")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n程序被用户中断")
    except Exception as e:
        print(f"\n程序执行出错: {e}")
        import traceback
        traceback.print_exc()
